import numpy as np
import matplotlib.pyplot as plt

class FussyLQRController:
    def __init__(self, A, B, Q, R):
        self.A = A
        self.B = B
        self.Q = Q
        self.R = R
        self.measurement_history = []
        self.setpoint_history = []

    def fuzzy_adjustment(self, error, error_dot):
        q_factor, r_factor = self.calculate_fuzzy_factors(error, error_dot)
        adjusted_Q = self.Q * q_factor
        adjusted_R = self.R * r_factor
        return adjusted_Q, adjusted_R

    def calculate_fuzzy_factors(self, error, error_dot):
        e_membership = self.gaussian_membership(error, [0, 1, 2], [0.5, 0.5, 0.5])
        e_dot_membership = self.gaussian_membership(error_dot, [0, 0.5, 1], [0.2, 0.2, 0.2])

        if e_membership[2] > 0.5 and e_dot_membership[2] > 0.5:
            q_factor = 1.0
            r_factor = 2.0
        elif e_membership[2] > 0.5 and e_dot_membership[1] > 0.5:
            q_factor = 1.0
            r_factor = 1.8
        elif e_membership[2] > 0.5 and e_dot_membership[0] > 0.5:
            q_factor = 1.0
            r_factor = 1.5
        elif e_membership[1] > 0.5 and e_dot_membership[2] > 0.5:
            q_factor = 1.5
            r_factor = 1.2
        elif e_membership[1] > 0.5 and e_dot_membership[1] > 0.5:
            q_factor = 1.2
            r_factor = 1.0
        elif e_membership[1] > 0.5 and e_dot_membership[0] > 0.5:
            q_factor = 1.0
            r_factor = 1.0
        elif e_membership[0] > 0.5 and e_dot_membership[2] > 0.5:
            q_factor = 0.8
            r_factor = 1.3
        elif e_membership[0] > 0.5 and e_dot_membership[1] > 0.5:
            q_factor = 0.7
            r_factor = 1.1
        elif e_membership[0] > 0.5 and e_dot_membership[0] > 0.5:
            q_factor = 0.6
            r_factor = 0.9
        elif e_membership[2] > 0.8 and e_dot_membership[0] > 0.8:
            q_factor = 1.1
            r_factor = 1.6
        elif e_membership[1] > 0.8 and e_dot_membership[1] > 0.8:
            q_factor = 1.3
            r_factor = 1.2
        elif e_membership[0] > 0.8 and e_dot_membership[2] > 0.8:
            q_factor = 0.7
            r_factor = 1.4
        else:
            q_factor = 0.5
            r_factor = 0.7

        return q_factor, r_factor

    def lqr_solve(self, Q=None, R=None):
        Q = Q if Q is not None else self.Q
        R = R if R is not None else self.R

        if np.isscalar(R):
            R = np.array([[R]])
        elif R.ndim == 1:
            R = np.diag(R)

        n = self.A.shape[0]
        m = self.B.shape[1]

        Z = np.vstack([
            np.hstack([self.A, -self.B @ np.linalg.inv(R) @ self.B.T]),
            np.hstack([-Q, -self.A.T])
        ])
        eigvals, eigvecs = np.linalg.eig(Z)

        stable_indices = np.where(np.real(eigvals) < 0)[0]
        eigvecs_stable = eigvecs[:, stable_indices]

        X = eigvecs_stable[:n, :]
        Y = eigvecs_stable[n:, :]

        P = np.real(Y @ np.linalg.inv(X))

        K = np.linalg.inv(R) @ self.B.T @ P

        eigVals, _ = np.linalg.eig(self.A - self.B @ K)

        return K, P, eigVals

    def update(self, setpoint, measurement,error_dot):
        self.measurement_history.append(measurement)
        self.setpoint_history.append(setpoint)
        error = setpoint - measurement
        if error_dot[0]==0:
            error_dot = np.zeros_like(error)
        adjusted_Q, adjusted_R = self.fuzzy_adjustment(error, error_dot)
        K, _, _ = self.lqr_solve(adjusted_Q, adjusted_R)
        u = K @ error
        return u

    def plot_error(self, title='none'):
        self.setpoint_history = np.array(self.setpoint_history)
        self.measurement_history = np.array(self.measurement_history)
        fig = plt.figure(figsize=(14, 4))
        subtitle = ['x', 'y', 'z']
        for i in range(len(self.setpoint_history[0])):
            ax = fig.add_subplot(131 + i)
            ax.plot([i / 100 for i in range(len(self.setpoint_history[:, i]))], self.setpoint_history[:, i],
                    label='setpoint')
            ax.plot([i / 100 for i in range(len(self.measurement_history[:, i]))], self.measurement_history[:, i],
                    label='measurement')
            ax.set_xlabel('time')
            ax.set_ylabel('value')
            ax.set_title(subtitle[i] + title)
            ax.legend()
        plt.tight_layout()

    def plot_fuzzy_rules(self):
        error_range = np.linspace(0, 2, 100)
        sigma = 0.5
        plt.figure(figsize=(12, 5))

        low = np.exp(-(error_range - 0)**2 / (2 * sigma**2))
        medium = np.exp(-(error_range - 1)**2 / (2 * sigma**2))
        high = np.exp(-(error_range - 2)**2 / (2 * sigma**2))

        plt.plot(error_range, low, label='Low Membership - Error')
        plt.plot(error_range, medium, label='Medium Membership - Error')
        plt.plot(error_range, high, label='High Membership - Error')

        plt.title('Error Membership Functions')
        plt.xlabel('Error Value')
        plt.ylabel('Membership Degree')
        plt.legend()
        plt.grid(True)

        error_dot_range = np.linspace(0, 1, 100)
        sigma = 0.2
        plt.figure(figsize=(12, 5))

        low_dot = np.exp(-(error_dot_range - 0)**2 / (2 * sigma**2))
        medium_dot = np.exp(-(error_dot_range - 0.5)**2 / (2 * sigma**2))
        high_dot = np.exp(-(error_dot_range - 1)**2 / (2 * sigma**2))

        plt.plot(error_dot_range, low_dot, label='Low Membership - Error_dot')
        plt.plot(error_dot_range, medium_dot, label='Medium Membership - Error_dot')
        plt.plot(error_dot_range, high_dot, label='High Membership - Error_dot')

        plt.title('Error_dot Membership Functions')
        plt.xlabel('Error_dot Value')
        plt.ylabel('Membership Degree')
        plt.legend()
        plt.grid(True)

    def gaussian_membership(self, value, centers, sigmas):
        membership = np.zeros(len(centers))

        for i, (center, sigma) in enumerate(zip(centers, sigmas)):
            membership[i] = np.exp(-(value[0] - center)**2 / (2 * sigma**2))

        return membership

    def plot_rule_matrix(self):
        error_range = np.linspace(0, 2, 100)
        error_dot_range = np.linspace(0, 1, 100)

        q_factors = np.zeros((len(error_range), len(error_dot_range)))
        r_factors = np.zeros((len(error_range), len(error_dot_range)))

        for i, error in enumerate(error_range):
            for j, error_dot in enumerate(error_dot_range):
                q_factors[i, j], r_factors[i, j] = self.calculate_fuzzy_factors(np.array([error]),
                                                                                np.array([error_dot]))

        fig, ax = plt.subplots(1, 2, figsize=(12, 6))

        c1 = ax[0].imshow(q_factors, extent=[0, 1, 0, 2], origin='lower', aspect='auto', cmap='viridis')
        ax[0].set_title('q_factor Matrix')
        ax[0].set_xlabel('Error_dot')
        ax[0].set_ylabel('Error')
        fig.colorbar(c1, ax=ax[0])

        c2 = ax[1].imshow(r_factors, extent=[0, 1, 0, 2], origin='lower', aspect='auto', cmap='viridis')
        ax[1].set_title('r_factor Matrix')
        ax[1].set_xlabel('Error_dot')
        ax[1].set_ylabel('Error')
        fig.colorbar(c2, ax=ax[1])

        plt.show()

if __name__ == '__main__':

    # 示例使用
    A = np.array([[0., 1.], [-1., -1.]])
    B = np.array([[1.], [1.]])
    Q = np.array([[1., 0.], [0., 1.]])
    R = np.array([[1.]])

    controller = FussyLQRController(A, B, Q, R)

    # 定义初始状态和设定点
    initial_state = np.array([10., 10.])
    setpoint = np.array([0., 0.])

    # 模拟系统响应
    time_steps = 100
    states = [initial_state]
    for t in range(time_steps):
        current_state = states[-1]
        control_input = controller.update(setpoint, current_state,0)
        print(control_input)
        new_state = (A @ current_state.reshape(-1,1) + B @ control_input.reshape(-1,1)).reshape(2)
        print(new_state)
        states.append(new_state)

    # 绘制误差图
    controller.plot_error(title="Position and Velocity Error")

    # 绘制模糊规则
    controller.plot_fuzzy_rules()
    controller.plot_rule_matrix()
    plt.show()